import sqlite3
import csv
from collections import defaultdict

# 数据库连接
conn = sqlite3.connect("E:\\immunedb\\immuno.db")
cursor = conn.cursor()

try:
    # 指定需要分析的HLA类型
    target_hlas = ["HLA-A*02:01", "HLA-A*01:01", "HLA-B*48:01"]
    print(f"分析的HLA类型：{target_hlas}")

    # 存储结果的字典：
    # result[hla] = {
    #     'counts': {tissue_count: peptide_count},  # 次数统计
    #     'tissues': {tissue_count: set(tissue_names)}  # 对应组织名称集合
    # }
    result = {
        hla: {
            'counts': defaultdict(int),
            'tissues': defaultdict(set)
        } for hla in target_hlas
    }

    for hla in target_hlas:
        print(f"\n处理HLA: {hla}")
        
        # 步骤1：获取该HLA中probability > 0.5的所有peptide_id
        cursor.execute("""
            SELECT DISTINCT peptide_id 
            FROM predictions 
            WHERE hla = ? AND probability > 0.5
        """, (hla,))
        peptide_ids = [row[0] for row in cursor.fetchall()]
        print(f"  找到{len(peptide_ids)}个符合条件的peptide_id")

        if not peptide_ids:
            continue

        # 步骤2：遍历每个peptide_id，统计组织出现次数及名称
        for idx, peptide_id in enumerate(peptide_ids, 1):
            # 通过peptide_id关联protein_id
            cursor.execute("""
                SELECT protein_id 
                FROM peptides 
                WHERE peptide_id = ?
            """, (peptide_id,))
            protein_id = cursor.fetchone()
            if not protein_id:
                continue  # 跳过无对应protein_id的肽段
            protein_id = protein_id[0]

            # 通过protein_id关联geneName
            cursor.execute("""
                SELECT geneName 
                FROM protein_gene 
                WHERE protein_id = ?
            """, (protein_id,))
            gene_name = cursor.fetchone()
            if not gene_name:
                continue  # 跳过无对应geneName的蛋白质
            gene_name = gene_name[0]

            # 通过geneName关联tissue（仅高表达，去重）
            cursor.execute("""
                SELECT DISTINCT tissue 
                FROM protein_exp 
                WHERE geneName = ? AND exp_level = 'High'
            """, (gene_name,))
            tissues = [row[0] for row in cursor.fetchall()]  # 提取组织名称列表
            tissue_count = len(tissues)

            # 更新统计结果：次数+1，记录组织名称
            if tissue_count > 0:
                result[hla]['counts'][tissue_count] += 1
                # 将该肽段对应的组织添加到集合（自动去重同一HLA下的重复组织）
                for tissue in tissues:
                    result[hla]['tissues'][tissue_count].add(tissue)
            
            # 进度提示（每100个肽段）
            if idx % 100 == 0:
                print(f"  已处理{idx}/{len(peptide_ids)}个肽段")

    # 步骤3：整理结果并写入CSV
    output_path = "E:/immunedb/hla_high_exp_tissue_details.csv"
    with open(output_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["HLA类型", "高表达组织出现次数", "肽段数量", "涉及的组织名称"])
        
        for hla in target_hlas:
            # 按组织出现次数升序排列
            tissue_counts = sorted(result[hla]['counts'].keys())
            for count in tissue_counts:
                # 组织名称用逗号拼接
                tissue_names = ",".join(sorted(result[hla]['tissues'][count]))
                writer.writerow([
                    hla,
                    count,
                    result[hla]['counts'][count],
                    tissue_names
                ])
            # 补充1-3次的记录（即使数量为0）
            for count in range(1, 4):
                if count not in result[hla]['counts']:
                    writer.writerow([hla, count, 0, ""])

    print(f"\n统计结果已保存到：{output_path}")
    print("\n统计摘要：")
    for hla in target_hlas:
        print(f"\n{HLA}：")
        for count in sorted(result[hla]['counts'].keys()):
            tissues = sorted(result[hla]['tissues'][count])
            print(f"  出现{count}次的肽段数量：{result[hla]['counts'][count]}，涉及组织：{tissues}")

except sqlite3.Error as e:
    print(f"\n数据库错误: {e}")
except Exception as e:
    print(f"\n其他错误: {e}")
finally:
    cursor.close()
    conn.close()
    print("\n程序结束，数据库连接已关闭")